package com.snail.starter.cloud.tracing;

import cn.hutool.extra.servlet.ServletUtil;
import com.snail.starter.utils.lambda.GenericEntityBuilder;
import com.snail.starter.cloud.annotation.IgnoreTracingLog;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.http.MediaType;
import org.springframework.web.filter.OncePerRequestFilter;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerExecutionChain;
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
import org.springframework.web.util.ContentCachingRequestWrapper;
import org.springframework.web.util.ContentCachingResponseWrapper;
import org.springframework.web.util.WebUtils;

import javax.annotation.Resource;
import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * @author: snail framework
 * @Date: 2023/2/22 下午5:06
 * @Version: 1.0.0
 */
@Slf4j
public class TracingLogFilter extends OncePerRequestFilter implements Ordered {
    private static final String FILE_BODY = "file body";

    private static final String LONG_BODY = "long body";

    @Override
    public int getOrder() {
        return Ordered.LOWEST_PRECEDENCE;
    }

    @Resource
    RequestMappingHandlerMapping mappingHandlerMapping;

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
        if (skip(request)) {
            filterChain.doFilter(request, response);
            return;
        }
        log.info("-----------------------Tracing Start---------------------------------");
        long start = System.currentTimeMillis();
        if (!(request instanceof ContentCachingRequestWrapper)) {
            request = new ContentCachingRequestWrapper(request);
        }
        if (!(response instanceof ContentCachingResponseWrapper)) {
            response = new ContentCachingResponseWrapper(response);
        }
        GenericEntityBuilder<Tracing> builder = GenericEntityBuilder.of(Tracing::new);
        try {
            filterChain.doFilter(request, response);
        } finally {
            builder.with(Tracing::setIp, MDC.get("ip"))
                    .with(Tracing::setPath, request.getRequestURI())
                    .with(Tracing::setMethod, request.getMethod())
                    .with(Tracing::setHeader, getHeaders(request))
                    .with(Tracing::setQueryParams, request.getQueryString())
                    .with(Tracing::setRequestBody, getRequestBody(request))
                    .with(Tracing::setResponseBody, getResponseBody(response))
                    .with(Tracing::setCostTime, System.currentTimeMillis() - start)
                    .with(Tracing::setStatus, response.getStatus());
            log.info("Tracing :\n {}", builder.build());
            log.info("-----------------------Tracing End-----------------------------------");
        }
    }

    private String getHeaders(HttpServletRequest request) {
        Enumeration<String> headerNames = request.getHeaderNames();
        Map<String, String> headerMap = new HashMap<>();
        while (headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            headerMap.put(headerName, request.getHeader(headerName));
        }
        return headerMap.toString();
    }

    private String getRequestBody(HttpServletRequest request) {
        if (ServletUtil.isMultipart(request)) {
            return FILE_BODY;
        }
        if (isFileContentType(request.getContentType())) {
            return FILE_BODY;
        }
        ContentCachingRequestWrapper wrapper = WebUtils.getNativeRequest(request, ContentCachingRequestWrapper.class);
        if (null == wrapper) {
            return null;
        }
        try {
            String body = new String(wrapper.getContentAsByteArray());
            if (null == body || body.length() > 20480) {
                return LONG_BODY;
            }
            return body;
        } catch (Throwable throwable) {
            log.warn("", throwable);
            return null;
        }
    }

    private String getResponseBody(HttpServletResponse response) {
        if (isFileContentType(response.getContentType())) {
            return FILE_BODY;
        }
        ContentCachingResponseWrapper wrapper = WebUtils.getNativeResponse(response, ContentCachingResponseWrapper.class);
        if (null == wrapper) {
            return null;
        }
        try {
            byte[] byteArray = wrapper.getContentAsByteArray();
            wrapper.copyBodyToResponse();
            String body = new String(byteArray);
            if (null == body || body.length() > 20480) {
                return LONG_BODY;
            }
            return body;
        } catch (Throwable throwable) {
            log.warn("", throwable);
            return null;
        }
    }

    private boolean isFileContentType(String contentType) {
        if (null == contentType) {
            return false;
        }
        MediaType mediaType = null;
        try {
            mediaType = MediaType.valueOf(contentType);
        } catch (Throwable throwable) {
            log.warn("", throwable);
        }
        if (null == mediaType) {
            return false;
        }
        if (MediaType.APPLICATION_OCTET_STREAM.equals(mediaType)) {
            return true;
        }
        if (MediaType.IMAGE_GIF.equals(mediaType)) {
            return true;
        }
        if (MediaType.IMAGE_JPEG.equals(mediaType)) {
            return true;
        }
        if (MediaType.IMAGE_PNG.equals(mediaType)) {
            return true;
        }
        return MediaType.APPLICATION_PDF.equals(mediaType);
    }

    private boolean skip(HttpServletRequest request) {
        HandlerExecutionChain handlerExecutionChain;
        try {
            handlerExecutionChain = mappingHandlerMapping.getHandler(request);
        } catch (Throwable throwable) {
            log.warn("mapping get handler error:{}", request.getRequestURI());
            return false;
        }
        if (null == handlerExecutionChain) {
            return false;
        }
        Object handler = handlerExecutionChain.getHandler();
        if (null == handler) {
            return false;
        }
        if (!(handler instanceof HandlerMethod)) {
            return false;
        }
        return ((HandlerMethod) handler).hasMethodAnnotation(IgnoreTracingLog.class);
    }
}